import threading
import serial
import cv2
import numpy as np
import time
from ultralytics import YOLO
from PID_control import pid_tracking_thread
from PID_control import car_control_callback

#小车左右轮速度
left_wheel_speed = 0
right_wheel_speed = 0

# 全局图像槽，每个线程取自己的那一份
images = [None, None, None]
conf_results = [None, None, None]

# 同步信号
frame_ready = [threading.Event() for _ in range(3)]  # 主线程投递新帧
done_event = [threading.Event() for _ in range(3)]   # 子线程处理完毕

def pedestrian_detect_thread_RGB(idx):
    #导入权重
    model = YOLO("RGB_pedestrian.pt")

    while True:
        # 等待主线程喂数据
        frame_ready[idx].wait()
        frame_ready[idx].clear()
        frame = images[idx]
        height, width, _ = frame.shape
        # 定义中心区域
        center_x1 = width // 3
        center_y1 = height // 3
        center_x2 = width * 2 // 3
        center_y2 = height * 2 // 3

        exist_pedestrian = False
        pedestrian_conf = None

        results = model(frame)
        boxes = results[0].boxes

        for box in boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            bbox_cx = (x1 + x2) // 2
            bbox_cy = (y1 + y2) // 2
            #如果中心区域出现行人，且置信度大于0.75,则向主线程发送消息
            if cls_id == 0 and conf > 0.75 \
                and center_x1 <= bbox_cx <= center_x2 \
                and center_y1 <= bbox_cy <= center_y2:
                
                exist_pedestrian = True
                pedestrian_conf = max(conf, pedestrian_conf)
        #传出结果
        if exist_pedestrian:
            conf_results[idx] = pedestrian_conf
        done_event[idx].set()  # 通知主线程已完成

def pedestrian_detect_thread_Thermal(idx):
    #导入权重
    model = YOLO("Thermal_pedestrian.pt")

    while True:
        # 等待主线程喂数据
        frame_ready[idx].wait()
        frame_ready[idx].clear()
        frame = images[idx]
        height, width, _ = frame.shape
        # 定义中心区域
        center_x1 = width // 3
        center_y1 = height // 3
        center_x2 = width * 2 // 3
        center_y2 = height * 2 // 3

        exist_pedestrian = False
        pedestrian_conf = None

        results = model(frame)
        boxes = results[0].boxes

        for box in boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            bbox_cx = (x1 + x2) // 2
            bbox_cy = (y1 + y2) // 2
            # 如果中心区域出现行人，且置信度大于0.75,则向主线程发送消息
            if cls_id == 0 and conf > 0.75 \
                and center_x1 <= bbox_cx <= center_x2 \
                and center_y1 <= bbox_cy <= center_y2:
                
                exist_pedestrian = True
                pedestrian_conf = max(conf, pedestrian_conf)

        # 传出结果
        if exist_pedestrian:
            conf_results[idx] = pedestrian_conf
        done_event[idx].set()  # 通知主线程已完成

def pedestrian_detect_thread_mid_fusion(idx):
    #导入权重
    model = YOLO("Mid_fusion_pedestrian.pt")

    while True:
        # 等待主线程喂数据
        frame_ready[idx].wait()
        frame_ready[idx].clear()
        frame = images[idx]
        height, width, _ = frame.shape
        # 定义中心区域
        center_x1 = width // 3
        center_y1 = height // 3
        center_x2 = width * 2 // 3
        center_y2 = height * 2 // 3

        exist_pedestrian = False
        pedestrian_conf = None

        results = model(frame)
        boxes = results[0].boxes

        for box in boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            bbox_cx = (x1 + x2) // 2
            bbox_cy = (y1 + y2) // 2
            # 如果中心区域出现行人，且置信度大于0.75,则向主线程发送消息
            if cls_id == 0 and conf > 0.75 \
                and center_x1 <= bbox_cx <= center_x2 \
                and center_y1 <= bbox_cy <= center_y2:
                
                exist_pedestrian = True
                pedestrian_conf = max(conf, pedestrian_conf)

        # 传出结果
        if exist_pedestrian:
            conf_results[idx] = pedestrian_conf
        done_event[idx].set()  # 通知主线程已完成
class PID:
    def __init__(self, Kp, Ki, Kd):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.last_error = 0
        self.integral = 0

    def compute(self, error, dt):
        self.integral += error * dt
        derivative = (error - self.last_error) / dt if dt > 0 else 0
        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        self.last_error = error
        return output
    
# 巡迹线程的回调，实时更新小车速度
def car_control_callback(adjust):
    global left_wheel_speed, right_wheel_speed
    base_speed = 50  # 可调
    max_speed = 100
    left_wheel_speed = int(max(min(base_speed - adjust, max_speed), 0))
    right_wheel_speed = int(max(min(base_speed + adjust, max_speed), 0))

def pid_tracking(pid_params, car_control_callback, frame):
    pid = PID(*pid_params)
    last_time = time.time()

    while True:
        # 图像预处理（如灰度、二值化）
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(gray, (5,5), 0)
        _, binary = cv2.threshold(blur, 100, 255, cv2.THRESH_BINARY_INV)

        # 选取ROI，只看下半部分
        height, width = binary.shape
        roi = binary[int(height*0.6):, :]

        # 寻找左右轨道
        scan_line = roi[-1, :]
        black_indices = np.where(scan_line == 0)[0]
        left_pos = black_indices[0] if len(black_indices) > 0 else 0
        right_pos = black_indices[-1] if len(black_indices) > 0 else width - 1

        # 4. 计算图像中线到轨道中点的偏移
        track_center = (left_pos + right_pos) // 2
        image_center = width // 2
        error = image_center - track_center

        now = time.time()
        dt = now - last_time
        last_time = now

        # 5. PID计算调整量（如输出给左右轮速差）
        adjust = pid.compute(error, dt)
        # 调用控制回调函数，调整小车速度/转向
        car_control_callback(adjust)
#PID类
class PID:
    def __init__(self, Kp, Ki, Kd):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.last_error = 0
        self.integral = 0

    def compute(self, error, dt):
        self.integral += error * dt
        derivative = (error - self.last_error) / dt if dt > 0 else 0
        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        self.last_error = error
        return output
    
# 巡迹线程的回调，实时更新小车速度
def car_control_callback(adjust):
    global left_wheel_speed, right_wheel_speed
    base_speed = 50  # 可调
    max_speed = 100
    left_wheel_speed = int(max(min(base_speed - adjust, max_speed), 0))
    right_wheel_speed = int(max(min(base_speed + adjust, max_speed), 0))

def pid_tracking_thread(pid_params, car_control_callback, frame):
    pid = PID(*pid_params)
    last_time = time.time()

    while True:
        # 图像预处理（如灰度、二值化）
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(gray, (5,5), 0)
        _, binary = cv2.threshold(blur, 100, 255, cv2.THRESH_BINARY_INV)

        # 选取ROI，只看下半部分
        height, width = binary.shape
        roi = binary[int(height*0.6):, :]

        # 寻找左右轨道
        scan_line = roi[-1, :]
        black_indices = np.where(scan_line == 0)[0]
        left_pos = black_indices[0] if len(black_indices) > 0 else 0
        right_pos = black_indices[-1] if len(black_indices) > 0 else width - 1

        # 4. 计算图像中线到轨道中点的偏移
        track_center = (left_pos + right_pos) // 2
        image_center = width // 2
        error = image_center - track_center

        now = time.time()
        dt = now - last_time
        last_time = now

        # 5. PID计算调整量（如输出给左右轮速差）
        adjust = pid.compute(error, dt)
        # 调用控制回调函数，调整小车速度/转向
        car_control_callback(adjust)

def main():
    #打开RGB摄像头
    rgb_cap = cv2.VideoCapture(0)
    if not rgb_cap.isOpened():
        print("无法打开RGB摄像头!")
        return
    #打开热成像摄像头
    thermal_cap = cv2.VideoCapture(1)
    if not thermal_cap.isOpened():
        print("无法打开热成像摄像头!")
        return   
    # 启动子线程
    threads = []
    threads.append(threading.Thread(target=pedestrian_detect_thread_RGB, args=(0,)))
    threads.append(threading.Thread(target=pedestrian_detect_thread_Thermal, args=(1,)))
    threads.append(threading.Thread(target=pedestrian_detect_thread_mid_fusion, args=(2,)))
    for t in threads:
        t.start()
    #循环获取帧
    while True:
        rgb_ret, rgb_frame = rgb_cap.read()
        if not rgb_ret:
            print("无法获取帧!")
            continue
        thermal_ret, thermal_frame = thermal_cap.read()
        if not thermal_ret:
            print("无法获取帧!")
            continue
        mid_fusion_frame = mix(rgb_frame, thermal_frame)
        # 捕获新帧
        images[0] = rgb_frame
        images[1] = thermal_frame
        images[2] = mid_fusion_frame
        # 清理上次done事件
        for e in done_event:
            e.clear()
        # 投递新帧，唤醒三个子线程
        for e in frame_ready:
            e.set()
        # 主线程等三个子线程都处理完
        for e in done_event:
            e.wait()
        fin_conf = bayesian_fusion(conf_results[0], conf_results[1], conf_results[2])
        pid_tracking([0.3, 0.01, 0.05], car_control_callback, images[0])
        # 决策：是否需要停车
        if fin_conf > 0.75:
            lw, rw = 0, 0
        else:
            lw, rw = left_wheel_speed, right_wheel_speed
        # 向小车主控发送指令
        cmd = f"L:{lw} R:{rw}\n"
        ser.write(cmd.encode('utf-8'))
    
    #关闭摄像头资源
    rgb_cap.release()
    thermal_cap.release()
    cv2.destroyAllWindows()
    #等待进程退出
    for t in threads:
        t.join()
    print("所有线程退出")

main()